/**
  ******************************************************************************
  * Copyright (c) 2022 - ~, SCUT-RobotLab Development Team
  * @file    FSM_State.cpp
  * @author
  * @brief   Code for .
  * @date   2022-02-28
  * @version 1.0
  * @par Change Log:
  * <table>
  * <tr><th>Date        <th>Version  <th>Author     <th>Description

  * </table>
  *
  ******************************************************************************
  * @attention
  *
  * if you had modified this file, please make sure your code does not have many
  * bugs, update the version Number, write dowm your name and the date, the most
  * important is make sure the users will have clear and definite understanding
  * through your new brief.
  *
  * <h2><center>&copy; Copyright (c) 2022 - ~,SCUT-RobotLab Development Team.
  * All rights reserved.</center></h2>
  ******************************************************************************
  */

/*============================= FSM State =============================*/
/**
 * FSM State base class
 */

#include <FSM_State.h>
#include <common/Utilities/Utilities_print.h>

/**
 * Constructor for the FSM State class.
 *
 * @param _controlFSMData holds all of the relevant control data
 * @param stateNameIn the enumerated state name
 * @param stateStringIn the string name of the current FSM state
 */
template <typename T>
FSM_State<T>::FSM_State(ControlFSMData<T>* _controlFSMData, FSM_StateName stateNameIn, std::string stateStringIn)
    : _data(_controlFSMData), stateName(stateNameIn), stateString(stateStringIn) {
    transitionData.zero();
    std::cout << "[FSM_State] Initialized FSM state: " << stateStringIn << std::endl;
}

/**
 * Cartesian impedance control for a given leg.
 *
 * @param leg the leg number to control
 * @param qDes desired joint position
 * @param dqDes desired joint velocity
 */
template <typename T>
void FSM_State<T>::jointPDControl(int leg, Vec3<T> qDes, Vec3<T> qdDes) {

    kpMat << 200, 0, 0, 0, 200, 0, 0, 0, 100;
    kdMat << 15, 0, 0, 0,15, 0, 0, 0, 20;

    _data->_legController->commands[leg].kpJoint = kpMat;
    _data->_legController->commands[leg].kdJoint = kdMat;

    _data->_legController->commands[leg].qDes = qDes;
    _data->_legController->commands[leg].qdDes = qdDes;
}

/**
 * Cartesian impedance control for a given leg.
 *
 * @param leg the leg number to control
 * @param pDes desired foot position
 * @param vDes desired foot velocity
 * @param kp_cartesian P gains
 * @param kd_cartesian D gains
 */
template <typename T>
void FSM_State<T>::cartesianImpedanceControl(int leg, Vec3<T> pDes, Vec3<T> vDes, Vec3<double> kp_cartesian,
                                             Vec3<double> kd_cartesian) {
    _data->_legController->commands[leg].pDes = pDes;
    // Create the cartesian P gain matrix
    kpMat << kp_cartesian[0], 0, 0, 0,
        kp_cartesian[1], 0, 0, 0,
        kp_cartesian[2];
    _data->_legController->commands[leg].kpCartesian = kpMat;

    _data->_legController->commands[leg].vDes = vDes;
    // Create the cartesian D gain matrix
    kdMat << kd_cartesian[0], 0, 0, 0, kd_cartesian[1], 0, 0, 0, kd_cartesian[2];
    _data->_legController->commands[leg].kdCartesian = kdMat;
    
    _data->_legController->setEnabled(leg, true);
}

/**
 *
 */
/*
template <typename T>
void FSM_State<T>::footstepHeuristicPlacement(int leg) {

  // Create the projection matrix for the 2D foot placement components
  Mat23<float> projectionMatrix;
  projectionMatrix << 1, 0, 0,
                   0, 1, 0;

  Vec3<float> velDes = _data->_desiredStateCommand->data.stateDes.block<3, 1>(6,
0); Vec3<float> angVelDes = _data->_desiredStateCommand->data.stateDes.block<3,
1>(9, 0); Mat3<float> rBody = _data->_stateEstimate.rBody;

  // Find each of the footstep locations for the swing feet
  for (int leg = 0; leg < 4; leg++) {
    if (_data->_gaitScheduler->gaitData.contactStateScheduled(leg)) {
      // The leg is in contact so nothing to do here

    } else {
      if (_data->_gaitScheduler->gaitData._currentGait ==
GaitType::TRANSITION_TO_STAND) {
        // Position the legs under the hips to stand...
        // Could also get rid of this and simply send 0 velocity ang vel
commands
        // from the CoM desired planner...
        Vec3<float> posHip = _data->_quadruped.getHipLocation(leg);
        footstepLocations.col(leg) <<
projectionMatrix.transpose()*projectionMatrix*
                                   (_data->_stateEstimate.position + // rBody *
posHip); } else {
        // Pull out the approximate yaw rate component of the robot in the
world. Vec3<float> yaw_rate; yaw_rate << 0, 0, _stateEstimate.omegaWorld(3);

        Vec3<float> posHip = _data->_quadruped.getHipLocation(leg);

        float timeStance = _data->_gaitScheduler->gaitData.timeStance(leg);

        // Footstep heuristic composed of several parts in the world frame
        footstepLocations.col(leg) <<
projectionMatrix.transpose()*projectionMatrix*      // Ground projection
                                   (_stateEstimate.position + // rBody * posHip
+                                      // Foot under hips timeStance / 2 *
velDes +                             // Raibert Heuristic timeStance / 2 *
(angVelDes.cross(rBody * posHip)) +  // Turning Raibert Heuristic
                                    (_stateEstimate.vBody - velDes));
      }
    }

  }
}
*/

/**
 * Gait independent formulation for choosing appropriate GRF and step locations
 * as well as converting them to leg controller understandable values.
 */
template <typename T>
void FSM_State<T>::runControls() {
    // This option should be set from the user interface or autonomously
    // eventually
    int CONTROLLER_OPTION = 1;

    // Reset the forces and steps to 0
    footFeedForwardForces = Mat34<T>::Zero();
    footstepLocations = Mat34<T>::Zero();

    // Choose the controller to run for picking step locations and balance forces
    if (CONTROLLER_OPTION == 0) {
        // Test to make sure we can control the robot
        // Test to make sure we can control the robot these will be calculated by
        // the controllers
        for (int leg = 0; leg < 4; leg++) {
            footFeedForwardForces.col(leg) << 0.0, 0.0, 0;    //-220.36;
                                                              // footFeedForwardForces.col(leg) = stateEstimate.rBody *
                                                              // footFeedForwardForces.col(leg);

            // footstepLocations.col(leg) << 0.0, 0.0,
            //     -_data->_quadruped->_maxLegLength / 2;
        }
    } else if (CONTROLLER_OPTION == 1) {
        // QP Balance Controller
        runBalanceController();

        // Swing foot landing positions are calculated with heuristics
        for (int leg = 0; leg < 4; leg++) {
            // footstepLocations.col(leg) << 0.0, 0.0,
            //     -_data->_quadruped->_maxLegLength / 2;
        }    // footstepHeuristicPlacement();

    } else if (CONTROLLER_OPTION == 2) {
        // WBC
        // runWholeBodyController();

    } else if (CONTROLLER_OPTION == 3) {
        // cMPC
        // runConvexModelPredictiveController();

        // Swing foot landing positions are calculated with heuristics
        // footstepHeuristicPlacement();

    } else if (CONTROLLER_OPTION == 4) {
        // RPC
        // runRegularizedPredictiveController();

    } else {
        // Zero out the commands if a controller was not selected
        jointFeedForwardTorques = Mat34<float>::Zero();    // feed forward joint torques
        jointPositions = Mat34<float>::Zero();             // joint angle positions
        jointVelocities = Mat34<float>::Zero();            // joint angular velocities
        footFeedForwardForces = Mat34<float>::Zero();      // feedforward forces at the feet
        footPositions = Mat34<float>::Zero();              // cartesian foot positions
        footVelocities = Mat34<float>::Zero();

        // Print an error message
        std::cout << "[FSM_State] ERROR: No known controller was selected: " << CONTROLLER_OPTION << std::endl;
    }
}

/**
 *
 */
template <typename T>
void FSM_State<T>::runBalanceController() {
    // double minForce = 25;
    // double maxForce = 500;
    // double contactStateScheduled[4];  // = {1, 1, 1, 1};
    // for (int i = 0; i < 4; i++) {
    //   contactStateScheduled[i] =
    //       _data->_gaitScheduler->gaitData.contactStateScheduled(i);
    // }

    // double minForces[4];  // = {minForce, minForce, minForce, minForce};
    // double maxForces[4];  // = {maxForce, maxForce, maxForce, maxForce};
    // for (int leg = 0; leg < 4; leg++) {
    //   minForces[leg] = contactStateScheduled[leg] * minForce;
    //   maxForces[leg] = contactStateScheduled[leg] * maxForce;
    // }

    // double COM_weights_stance[3] = {1, 1, 10};
    // double Base_weights_stance[3] = {20, 10, 10};
    // double pFeet[12], p_des[3], p_act[3], v_des[3], v_act[3], O_err[3], rpy[3],
    //     omegaDes[3];
    // double se_xfb[13];
    // double kpCOM[3], kdCOM[3], kpBase[3], kdBase[3];

    // for (int i = 0; i < 4; i++) {
    //   se_xfb[i] = (double)_data->_stateEstimator->getResult().orientation(i);
    // }
    // // se_xfb[3] = 1.0;
    // for (int i = 0; i < 3; i++) {
    //   rpy[i] = 0;  //(double)_data->_stateEstimator->getResult().rpy(i);
    //   p_des[i] = (double)_data->_stateEstimator->getResult().position(i);
    //   p_act[i] = (double)_data->_stateEstimator->getResult().position(i);
    //   omegaDes[i] =
    //       0;  //(double)_data->_stateEstimator->getResult().omegaBody(i);
    //   v_act[i] = (double)_data->_stateEstimator->getResult().vBody(i);
    //   v_des[i] = (double)_data->_stateEstimator->getResult().vBody(i);

    //   se_xfb[4 + i] = (double)_data->_stateEstimator->getResult().position(i);
    //   se_xfb[7 + i] = (double)_data->_stateEstimator->getResult().omegaBody(i);
    //   se_xfb[10 + i] = (double)_data->_stateEstimator->getResult().vBody(i);

    //   // Set the translational and orientation gains
    //   kpCOM[i] = (double)_data->controlParameters->kpCOM(i);
    //   kdCOM[i] = (double)_data->controlParameters->kdCOM(i);
    //   kpBase[i] = (double)_data->controlParameters->kpBase(i);
    //   kdBase[i] = (double)_data->controlParameters->kdBase(i);
    // }

    // Vec3<T> pFeetVec;
    // Vec3<T> pFeetVecCOM;
    // // Get the foot locations relative to COM
    // for (int leg = 0; leg < 4; leg++) {
    //   computeLegJacobianAndPosition(**&_data->_quadruped,
    //                                 _data->_legController->datas[leg].q,
    //                                 (Mat3<T>*)nullptr, &pFeetVec, 1);
    //   //pFeetVecCOM = _data->_stateEstimator->getResult().rBody.transpose() *
    //                 //(_data->_quadruped->getHipLocation(leg) + pFeetVec);

    //   pFeetVecCOM = _data->_stateEstimator->getResult().rBody.transpose() *
    //                 (_data->_quadruped->getHipLocation(leg) + _data->_legController->datas[leg].p);

    //   pFeet[leg * 3] = (double)pFeetVecCOM[0];
    //   pFeet[leg * 3 + 1] = (double)pFeetVecCOM[1];
    //   pFeet[leg * 3 + 2] = (double)pFeetVecCOM[2];
    // }

    // balanceController.set_alpha_control(0.01);
    // balanceController.set_friction(0.5);
    // balanceController.set_mass(46.0);
    // balanceController.set_wrench_weights(COM_weights_stance, Base_weights_stance);
    // balanceController.set_PDgains(kpCOM, kdCOM, kpBase, kdBase);
    // balanceController.set_desiredTrajectoryData(rpy, p_des, omegaDes, v_des);
    // balanceController.SetContactData(contactStateScheduled, minForces, maxForces);
    // balanceController.updateProblemData(se_xfb, pFeet, p_des, p_act, v_des, v_act,
    //                                     O_err, 0.0);

    // double fOpt[12];
    // balanceController.solveQP_nonThreaded(fOpt);

    // // Publish the results over LCM
    // balanceController.publish_data_lcm();

    // // Copy the results to the feed forward forces
    // for (int leg = 0; leg < 4; leg++) {
    //   footFeedForwardForces.col(leg) << (T)fOpt[leg * 3], (T)fOpt[leg * 3 + 1],
    //       (T)fOpt[leg * 3 + 2];
    // }
}

/**
 * Gait independent formulation for choosing appropriate GRF and step locations
 * as well as converting them to leg controller understandable values.
 */
template <typename T>
void FSM_State<T>::turnOnAllSafetyChecks() {
    // Pre controls safety checks
    checkSafeOrientation = true;    // check roll and pitch

    // Post control safety checks
    checkPDesFoot = true;            // do not command footsetps too far
    checkForceFeedForward = true;    // do not command huge forces
    checkLegSingularity = true;      // do not let leg
}

/**
 *
 */
template <typename T>
void FSM_State<T>::turnOffAllSafetyChecks() {
    // Pre controls safety checks
    checkSafeOrientation = false;    // check roll and pitch

    // Post control safety checks
    checkPDesFoot = false;            // do not command footsetps too far
    checkForceFeedForward = false;    // do not command huge forces
    checkLegSingularity = false;      // do not let leg
}

// template class FSM_State<double>;
template class FSM_State<float>;

/************************ COPYRIGHT(C) SCUT-RobotLab Development Team **************************/
